from enum import Enum
from typing import Dict


class StompCommand(str, Enum):
    CONNECT = "CONNECT"
    SEND = "SEND"
    SUBSCRIBE = "SUBSCRIBE"
    UNSUBSCRIBE = "UNSUBSCRIBE"
    DISCONNECT = "DISCONNECT"


class StompFrame:
    def __init__(self, command: StompCommand, headers: Dict[str, str], body: str = ""):
        self.command = command
        self.headers = headers
        self.body = body

    @classmethod
    def parse(cls, raw: str) -> "StompFrame":
        lines = raw.strip().split("\n")
        command = StompCommand(lines[0])
        headers = {}
        i = 1
        while i < len(lines) and lines[i] != "":
            key, value = lines[i].split(":", 1)
            headers[key] = value
            i += 1
        body = "\n".join(lines[i + 1 :]) if i + 1 < len(lines) else ""
        return cls(command, headers, body)

    def to_frame(self) -> str:
        header_lines = [f"{k}:{v}" for k, v in self.headers.items()]
        return "\n".join([self.command.value] + header_lines + ["", self.body, "\x00"])
